
import sys
import indexing.featureindex

# CODON_INFO_PATH = "/home/bowtie/Bowtie_Illumina_Alignment/references/indexed_human_genome"
CODON_INFO_PATH = "/archive10/bowtie/references/ucsc_codon_info"
DBSNP_FILE = "/archive10/bowtie/references/known_snps/snp130-singles.txt"
SHENDURE_SNP_FILE = "/archive10/bowtie/references/known_snps/shendure_snps.txt"

resComp = { "A":"T", "T":"A", "C":"G", "G":"C", "N":"N", "a":"t", "t":"a", "g":"c", "c":"g", "n":"n" }

def toStrand(strand, base):
    if strand=="+": return base
    elif strand=="-": return resComp[base]
    raise Exception("bad strand: " + strand)
    
def readDbsnpSnps(filename=DBSNP_FILE):
    dbsnp = {}
    f = open(filename, "r")
    for line in f:
        cols = line.strip().split("\t")
        if len(cols)!=3:
            print >> sys.stderr, "# Warning: ignoring line", line.strip()
            continue
        key = (cols[0], cols[1], cols[2])
        dbsnp[key] = True
        if 0 == len(dbsnp)%100000:
            print >> sys.stderr, "# Read %d lines from %s" % (len(dbsnp), filename)
    f.close()
    print >> sys.stderr, "# Done... read %d lines from %s" % (len(dbsnp), filename)
    return dbsnp

def readBlosumMatrix(codonTable, filename="blosum62.txt"):
    codons = {}
    for (pattern, (aa, aaChar)) in codonTable.items():
        codons[aaChar] = aa
    rows = []
    mat = {}
    for line in open(filename):
        if line.startswith("#"):
            continue
        entries = filter(lambda(x):x!='', line.strip().split(" "))
        if not rows:
            rows = [codons.get(entry, entry) for entry in entries]
            print >> sys.stderr, "# Read BLOSUM column headers: ", rows
            for row in rows:
                mat[row] = {}
            continue
        if len(entries) != len(rows)+1:
            print >> sys.stderr, "# Ignoring BLOSUM file line: ", line.strip()
            continue
        row = codons.get(entries[0], entries[0])
        for ix in xrange(1, len(entries)):
            mat[row][rows[ix-1]] = int(entries[ix])
    print >> sys.stderr, "# BLOSUM matrix: ", mat
    return mat

def readCodonTable(filename="codon-table.txt"):
    table = {}
    f = open(filename, "r")
    for line in f:
        cols = line.strip().split("\t")
        if len(cols)!=3 or len(cols[0])!=3: continue
        (codon, aa, aaChar) = (cols[0], cols[1], cols[2])
        if codon[-1]=='*':
            codons = [codon[:2]+ch for ch in ["A", "C", "G", "T"]]
        else:
            codons = [codon]
        for codon in codons:
            table[codon] = (aa, aaChar)
    f.close()
    return table

def readGeneSymbols(filename="kgXref.table.txt"):
    mapping = {}
    f = open(filename, "r")
    for line in f:
        cols = line.strip().split("\t")
        mapping[cols[0]] = "(None)"
        if cols[4]!="": mapping[cols[0]] = cols[4]
    f.close()
    return mapping

def buildGeneLookup(filename="knownGene.table.txt"):
    lookup = {}
    symbols = readGeneSymbols()
    f = open(filename, "r")
    for line in f:
        if line.startswith("#"): continue
        cols = line.strip().split("\t")
        if not lookup.has_key(cols[1]): lookup[cols[1]] = []
        # (txStart, txEnd, ucscId, geneSymbol or "(None)", cdsStart, cdsEnd, strand)
        lookup[cols[1]].append((cols[3], cols[4], cols[0], symbols[cols[0]], long(cols[5]), long(cols[6]), cols[2]))
    f.close()
    for key in lookup:
        lookup[key].sort()
        print >> sys.stderr, "# %d entries in lookup table for %s" % (len(lookup[key]), key)
    return lookup

def getGeneSymbol(chr, pos, lookup):
    lookupChr = lookup[chr]
    for (spos, epos, ucscName, geneSymbol, cdsStart, cdsEnd, geneStrand) in lookupChr:
        if pos>=long(spos) and pos<=long(epos): return geneSymbol
    return "(None)"

def getGeneIds(chr, pos, lookup):
    ids = []
    lookupChr = lookup[chr]
    for (spos, epos, ucscName, geneSymbol, cdsStart, cdsEnd, geneStrand) in lookupChr:
        if pos>=long(spos) and pos<=long(epos):
            ids.append(ucscName)
    return ids

def getNonCodingLocationDesc(chr, pos, ucscId, lookup):
    lookupChr = lookup[chr]
    for (spos, epos, ucscName, geneSymbol, cdsStart, cdsEnd, geneStrand) in lookupChr:
        if ucscName==ucscId:
            if cdsStart==cdsEnd:
                return (geneStrand, "non-coding-gene")
            elif (geneStrand=="+" and pos<cdsStart) or (geneStrand=="-" and pos>cdsEnd):
                return (geneStrand, "non-coding-5-prime-utr")
            elif (geneStrand=="+" and pos>cdsEnd) or (geneStrand=="-" and pos<cdsStart):
                return (geneStrand, "non-coding-3-prime-utr")
            else:
                return (geneStrand, "non-coding-intronic")
    raise Exception("couldn't find " + ucscId + " in lookup table.")

def getCodons(positions):
    ret = positions
    for key in ret.keys():
        ret[key] = []
    print >> sys.stderr, "# Getting codons for %d positions..." % len(ret)
    for filename in [CODON_INFO_PATH + "/chr%s.ucsc-codon-info.txt"%x for x in (range(1,23)+["X","Y"])]:
    # for filename in [CODON_INFO_PATH + "/chr1.ucsc-codon-info.txt"]:
        print >> sys.stderr, "# ... parsing %s ..." % filename
        hits = 0
        f = open(filename, "r")
        for line in f:
            cols = line.strip().split("\t")
            (name, chrom, strand, pos, codonNum, codon, frame) = tuple(cols)
            key = (chrom, long(pos)+1)
            if ret.has_key(key):
                hits += 1
                ret[key].append((name, strand, codonNum, codon, frame))
                # print key, " --> ", ret[key]
        f.close()
        print >> sys.stderr, "# ... %d hits." % hits
    return ret

if __name__=="__main__":
    if len(sys.argv)<=1:
        print "Usage: python %s SNP-COUNTS" % sys.argv[0]
        sys.exit(1)
    codonTable = readCodonTable()
    blosum = readBlosumMatrix(codonTable)
    shensnps = readDbsnpSnps(SHENDURE_SNP_FILE)
    dbsnps = readDbsnpSnps()
    snps = {}  # entries like "chr6_234524_A_T" --> tab-separated columns containing SNP id
    positions = {}  # entries like ("chr1", 1414) --> True
    rowContents = {}   # entries like "chr6_234524_A_T" --> tab-separated columns with add'l SNP info
    for line in open(sys.argv[1]):
        cols = line.strip().split("\t")
        if len(cols)>0 and cols[0].startswith("chr") and len(cols[0].split(":"))==2:
            # first column is chr:pos, second column is old base, third column is new base
            (chr, pos) = tuple(cols[0].split(":"))
            snpId = "%s_%s_%s_%s" % (chr, pos, cols[1], cols[2])
            positions[(chr, long(pos))] = True
            snps[snpId] = "%s:%s\t%s\t%s" % (chr, pos, cols[1], cols[2])
            rowContents[snpId] = "\t".join(cols[3:])
        elif len(cols)==13:
            # G 19 0.7600 chr1 4793 True A 30 0.764474 20 T 6/0/24/0 6/0/19/0
            snpId = "%s_%s_%s_%s" % (cols[3], cols[4], cols[6], cols[0])
            snps[snpId] = line.strip()
            positions[(cols[3], long(cols[4]))] = True
            rowContents[snpId] = ""
        elif len(cols)==22:
            # reading from snp-results.txt file:
            if cols[0]=="Sample": continue
            snps["chr%s_%s_%s_%s" % (cols[2], cols[3], cols[4], cols[5])] = line.strip()
            positions[("chr%s" % cols[2], long(cols[3]))] = True
        elif len(cols)>6:
            # arbitrary columns where 2nd column is SNP id (dropping first column)
            snps[cols[1]] = "\t".join(cols[1].replace("chr", "").split("_"))
            fields = cols[1].split("_")
            positions[(fields[0], long(fields[1]))] = True
            rowContents[cols[1]] = "\t".join(cols[2:])
        else:
            snps["%s_%s_%s_%s" % (cols[0], cols[1], cols[2], cols[3])] = "%s\t%s\t%s\t%s" % (cols[0].replace("chr", ""), cols[1], cols[2], cols[3])
            positions[(cols[0], long(cols[1]))] = True
        if 0==len(positions)%10000:
            print >> sys.stderr, "# Read %d positions from file ... " % len(positions)
    annots = getCodons(positions)
    geneLookup = buildGeneLookup()
    for snp in snps:
        fields = snp.split("_")
        (chr, pos, oldBase, newBase) = (fields[0], long(fields[1]), fields[2], fields[3])
        dbtxt = "dbsnp-no"
        if dbsnps.has_key(("%s:%d" % (chr, pos), oldBase, newBase)):
            dbtxt = "dbsnp-yes"
        #if getDbsnp((chr, pos-1, pos+1)).has_key((chr, pos)):
        #    dbtxt = "dbsnp-yes"
        shentxt = "shendure-snp-no"
        if shensnps.has_key(("%s:%d" % (chr, pos), oldBase, newBase)):
            shentxt = "shendure-snp-yes"
        annot = annots[(chr, pos)]
        ucscNames = getGeneIds(chr, pos, geneLookup)
        symbol = getGeneSymbol(chr, pos, geneLookup)
        foundAnnotation = False
        annotatedNames = {}
        for (ucscName, strand, codonNum, codon, frame) in annot:
            annotatedNames[ucscName] = True
        if len(ucscNames)==0:
            # if no associated transcripts at all, print a single row
            desc = "non-coding-unannotated"
            print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (symbol, snp, snps[snp], rowContents[snp], dbtxt, shentxt, "non-coding-unannotated", "\t".join(["NA"]*8))
            # print "%s\t%s\t%s\t%s\t%s" % (symbol, snp, snps[snp], rowContents[snp], "non-coding-unannotated")
            continue
        noncodingresults = {}
        for ucscName in ucscNames:
            # there are associated transcripts, combine those with same description of location
            if annotatedNames.has_key(ucscName):
                continue
            (strand, desc) = getNonCodingLocationDesc(chr, pos, ucscName, geneLookup)
            key = (strand, desc)
            if noncodingresults.has_key(key): noncodingresults[key] += ","
            noncodingresults[key] = noncodingresults.get(key, "") + ucscName
            if desc!="non-coding-gene":
                foundAnnotation = True
        results = {}
        for (ucscName, strand, codonNum, codon, frame) in annot:
            foundAnnotation = True
            key = (strand, codon, int(frame))
            if results.has_key(key): results[key] += ","
            results[key] = results.get(key, "") + ucscName
        for (strand, codon, frame) in results:
            if codon[frame] != toStrand(strand, oldBase):
                raise Exception("inconsistent old base at " + snp)
            newCodon = ""
            for i in [0,1,2]:
                if i!=frame: newCodon += codon[i]
                else: newCodon += toStrand(strand, newBase)
            oldAmino = codonTable[codon][0]
            newAmino = codonTable[newCodon][0]
            type = "syn"
            if oldAmino != newAmino:
                type = "non-syn"
            print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%d\t%s" % (symbol, snp, snps[snp], rowContents[snp], dbtxt, shentxt, type, strand, frame+1, codon, newCodon, oldAmino, newAmino, blosum[oldAmino][newAmino], results[(strand, codon, frame)])
        for (strand, desc) in noncodingresults:
            if not foundAnnotation or desc!="non-coding-gene":
                print "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s" % (symbol, snp, snps[snp], rowContents[snp], dbtxt, shentxt, desc, strand, "NA", "NA", "NA", "NA", "NA", "NA", noncodingresults[(strand, desc)])
